# video_inference_4class_complete.py
import cv2
import numpy as np
from openvino.runtime import Core
import os
import time
import warnings
import logging
import subprocess
import tempfile
import shutil

# -----------------------------
# 配置路径（请根据你的实际路径修改！）
# -----------------------------
VEHICLE_MODEL_PATH = r"D:\CodeCNN\intel\vehicle-detection-0202\FP32\vehicle-detection-0202.xml"
ROAD_MODEL_PATH = r"D:\CodeCNN\intel\road-segmentation-adas-0001\FP32\road-segmentation-adas-0001.xml"
INPUT_FOLDER = r"H:\xiaomi\test"  # 输入图片文件夹
OUTPUT_FOLDER = r"detect"  # 输出结果文件夹
VIOLATION_FOLDER = r"violations"  # 压线违规图片文件夹
VIDEO_INPUT = r"H:\xiaomi\test\video\192.168.1.64_01_20251027093805446.MP4"  # 视频文件路径

# 创建输出目录
os.makedirs(OUTPUT_FOLDER, exist_ok=True)
os.makedirs(VIOLATION_FOLDER, exist_ok=True)

# -----------------------------
# 配置参数
# -----------------------------
DETECTION_INTERVAL = 3  # 视频检测间隔帧数
PREPROCESS_METHOD = 'adaptive'  # 预处理方法
CONFIDENCE_THRESHOLD = 0.001  # 压线检测阈值
SIDE_MARGIN_RATIO = 0.2  # 侧边边界比例

# 设置环境变量
os.environ['GLOG_minloglevel'] = '2'
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
warnings.filterwarnings('ignore')

# -----------------------------
# 初始化 OpenVINO
# -----------------------------
core = Core()
device = "CPU"  # 可选 "GPU"（若支持）

# -----------------------------
# 1. 加载车辆检测模型
# -----------------------------
print("Loading vehicle detection model...")
vehicle_model = core.read_model(model=VEHICLE_MODEL_PATH)
vehicle_compiled = core.compile_model(model=vehicle_model, device_name=device)
vehicle_input_layer = vehicle_compiled.input(0)
vehicle_output_layer = vehicle_compiled.output(0)

# 获取输入尺寸
_, _, h_det, w_det = vehicle_input_layer.shape
print(f"Vehicle detection input shape: {h_det}x{w_det}")

# -----------------------------
# 2. 加载道路分割模型（4类）
# -----------------------------
print("Loading 4-class road segmentation model...")
road_model = core.read_model(model=ROAD_MODEL_PATH)
road_compiled = core.compile_model(model=road_model, device_name=device)
road_input_layer = road_compiled.input(0)
road_output_layer = road_compiled.output(0)

# 输出形状应为 [1, 4, H, W]
_, num_classes, h_seg, w_seg = road_output_layer.shape
print(f"Road segmentation output classes: {num_classes} (should be 4)")
if num_classes != 4:
    print("⚠️ 模型输出不是4类！请确认是否使用了多分类版本（如 -1024x768）")
    exit(1)

print(f"Road segmentation input shape: {h_seg}x{w_seg}")

# -----------------------------
# 3. 定义 4 类颜色映射（BGR 格式）- 背景透明
# -----------------------------
COLOR_MAP = {
    0: None,  # 背景 - 透明
    1: (0, 255, 0),  # green - 道路
    2: (255, 0, 0),  # blue - 路缘
    3: (0, 0, 255),  # red - 道路标记
}


# -----------------------------
# 4. 图像预处理增强函数
# -----------------------------
def preprocess_image_for_segmentation(image, method='adaptive'):
    """
    对输入图像进行预处理，增强对比度以改善分割效果
    """
    processed = image.copy()

    if method == 'adaptive' or method == 'all':
        # 自适应方法：根据图像亮度选择最佳预处理
        gray = cv2.cvtColor(processed, cv2.COLOR_BGR2GRAY)
        brightness = np.mean(gray)

        if brightness < 80:  # 暗图像
            print(f"    🌙 检测到暗图像 (亮度: {brightness:.1f})，使用组合增强")
            # Gamma校正提亮
            gamma = 1.8
            inv_gamma = 1.0 / gamma
            table = np.array([((i / 255.0) ** inv_gamma) * 255 for i in np.arange(0, 256)]).astype("uint8")
            processed = cv2.LUT(processed, table)

            # CLAHE增强对比度
            lab = cv2.cvtColor(processed, cv2.COLOR_BGR2LAB)
            lab_planes = list(cv2.split(lab))
            clahe = cv2.createCLAHE(clipLimit=3.0, tileGridSize=(8, 8))
            lab_planes[0] = clahe.apply(lab_planes[0])
            lab = cv2.merge(lab_planes)
            processed = cv2.cvtColor(lab, cv2.COLOR_LAB2BGR)

        elif brightness > 200:  # 过亮图像
            print(f"    ☀️ 检测到过亮图像 (亮度: {brightness:.1f})，使用对比度增强")
            # 过亮图像主要增强对比度
            lab = cv2.cvtColor(processed, cv2.COLOR_BGR2LAB)
            lab_planes = list(cv2.split(lab))
            clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
            lab_planes[0] = clahe.apply(lab_planes[0])
            lab = cv2.merge(lab_planes)
            processed = cv2.cvtColor(lab, cv2.COLOR_LAB2BGR)

        else:  # 正常亮度
            print(f"    ⛅ 正常亮度图像 (亮度: {brightness:.1f})，使用标准增强")
            # 正常图像使用标准CLAHE
            lab = cv2.cvtColor(processed, cv2.COLOR_BGR2LAB)
            lab_planes = list(cv2.split(lab))
            clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
            lab_planes[0] = clahe.apply(lab_planes[0])
            lab = cv2.merge(lab_planes)
            processed = cv2.cvtColor(lab, cv2.COLOR_LAB2BGR)

    elif method == 'clahe':
        # CLAHE 对比度受限自适应直方图均衡化
        print("    🔧 使用CLAHE增强")
        lab = cv2.cvtColor(processed, cv2.COLOR_BGR2LAB)
        lab_planes = list(cv2.split(lab))
        clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
        lab_planes[0] = clahe.apply(lab_planes[0])
        lab = cv2.merge(lab_planes)
        processed = cv2.cvtColor(lab, cv2.COLOR_LAB2BGR)

    elif method == 'histeq':
        # 直方图均衡化
        print("    🔧 使用直方图均衡化")
        ycrcb = cv2.cvtColor(processed, cv2.COLOR_BGR2YCrCb)
        ycrcb[:, :, 0] = cv2.equalizeHist(ycrcb[:, :, 0])
        processed = cv2.cvtColor(ycrcb, cv2.COLOR_YCrCb2BGR)

    elif method == 'gamma':
        # Gamma 校正
        gamma = 1.5
        print(f"    🔧 使用Gamma校正 (gamma={gamma})")
        inv_gamma = 1.0 / gamma
        table = np.array([((i / 255.0) ** inv_gamma) * 255 for i in np.arange(0, 256)]).astype("uint8")
        processed = cv2.LUT(processed, table)

    return processed


# -----------------------------
# 5. 增强版道路分割推理
# -----------------------------
def enhanced_road_segmentation(image, preprocess_method='adaptive'):
    """
    增强版道路分割，包含图像预处理
    """
    # 图像预处理
    print(f"    应用图像预处理: {preprocess_method}")
    processed_image = preprocess_image_for_segmentation(image, preprocess_method)

    # 调整尺寸并进行推理
    resized_seg = cv2.resize(processed_image, (w_seg, h_seg))
    input_seg = np.expand_dims(resized_seg.transpose(2, 0, 1), axis=0).astype(np.float32)

    # 推理
    infer_result = road_compiled([input_seg])
    tensor_output = infer_result[road_output_layer]
    seg_probs = np.array(tensor_output)
    seg_class = np.argmax(seg_probs[0], axis=0).astype(np.uint8)

    # 缩放回原图尺寸
    orig_h, orig_w = image.shape[:2]
    seg_class_resized = cv2.resize(seg_class, (orig_w, orig_h), interpolation=cv2.INTER_NEAREST)

    return seg_class_resized


# -----------------------------
# 6. 优化版车辆压线检测函数
# -----------------------------
def check_lane_violation_optimized(vehicle_bbox, lane_marking_mask, confidence_threshold=0.01, side_margin_ratio=0.2):
    """
    优化版压线检测 - 只检测车辆底部中心区域，避开两侧
    """
    x_min, y_min, x_max, y_max = vehicle_bbox

    # 提取车辆底部区域
    bottom_height_ratio = 0.15
    bottom_y_min = int(y_max - (y_max - y_min) * bottom_height_ratio)

    # 排除两侧区域，只检测中心部分
    side_margin = int((x_max - x_min) * side_margin_ratio)
    center_x_min = x_min + side_margin
    center_x_max = x_max - side_margin

    # 确保坐标在图像范围内
    bottom_y_min = max(bottom_y_min, 0)
    bottom_y_max = min(y_max, lane_marking_mask.shape[0])
    center_x_min = max(center_x_min, 0)
    center_x_max = min(center_x_max, lane_marking_mask.shape[1])

    # 如果中心区域太小，使用完整宽度
    if center_x_max <= center_x_min or (center_x_max - center_x_min) < 10:
        center_x_min = x_min
        center_x_max = x_max
        print("    ⚠️ 中心区域过小，使用完整宽度检测")

    vehicle_bottom_bbox = [center_x_min, bottom_y_min, center_x_max, bottom_y_max]

    # 创建车辆底部中心区域的掩码
    vehicle_bottom_mask = np.zeros_like(lane_marking_mask)
    if bottom_y_max > bottom_y_min and center_x_max > center_x_min:
        vehicle_bottom_mask[bottom_y_min:bottom_y_max, center_x_min:center_x_max] = 1

    # 计算重叠区域
    overlap_mask = vehicle_bottom_mask & lane_marking_mask
    overlap_pixels = np.sum(overlap_mask)

    # 计算重叠比例
    vehicle_bottom_area = (center_x_max - center_x_min) * (bottom_y_max - bottom_y_min)
    if vehicle_bottom_area > 0:
        overlap_ratio = overlap_pixels / vehicle_bottom_area
    else:
        overlap_ratio = 0

    # 判断是否压线
    is_violation = overlap_ratio > confidence_threshold

    return is_violation, overlap_ratio, overlap_pixels, vehicle_bottom_bbox, overlap_mask


# -----------------------------
# 7. 保存压线违规图片函数
# -----------------------------
def save_violation_image(image, filename, vehicle_bbox, bottom_bbox, overlap_mask,
                         violation_count, total_vehicles, overlap_ratio, is_video=False, frame_num=0):
    """
    保存压线违规图片，包含详细标注信息
    """
    # 创建违规图片副本
    violation_img = image.copy()

    # 绘制车辆边界框（红色）
    x_min, y_min, x_max, y_max = vehicle_bbox
    cv2.rectangle(violation_img, (x_min, y_min), (x_max, y_max), (0, 0, 255), 3)

    # 绘制检测底部区域（黄色）
    bx_min, by_min, bx_max, by_max = bottom_bbox
    cv2.rectangle(violation_img, (bx_min, by_min), (bx_max, by_max), (0, 255, 255), 2)

    # 绘制重叠区域（红色高亮）
    overlap_coords = np.where(overlap_mask)
    for y, x in zip(overlap_coords[0], overlap_coords[1]):
        cv2.circle(violation_img, (x, y), 2, (0, 0, 255), -1)

    # 添加详细的标注信息
    if is_video:
        info_lines = [
            f"VIDEO VIOLATION DETECTED!",
            f"Frame: {frame_num}",
            f"Overlap Ratio: {overlap_ratio:.3f}",
            f"Vehicle BBox: [{x_min}, {y_min}, {x_max}, {y_max}]",
            f"Time: {time.strftime('%H:%M:%S')}",
            f"Violations: {violation_count}"
        ]
    else:
        info_lines = [
            f"VIOLATION DETECTED!",
            f"Overlap Ratio: {overlap_ratio:.3f}",
            f"Vehicle BBox: [{x_min}, {y_min}, {x_max}, {y_max}]",
            f"Total Vehicles: {total_vehicles}",
            f"Violations: {violation_count}"
        ]

    # 绘制信息背景
    text_bg_height = len(info_lines) * 25 + 20
    cv2.rectangle(violation_img, (10, 10), (400, text_bg_height), (0, 0, 0), -1)
    cv2.rectangle(violation_img, (10, 10), (400, text_bg_height), (0, 0, 255), 2)

    # 绘制信息文本
    for i, line in enumerate(info_lines):
        y_pos = 40 + i * 25
        color = (0, 0, 255) if i == 0 else (255, 255, 255)
        font_scale = 0.7 if i == 0 else 0.5
        thickness = 2 if i == 0 else 1
        cv2.putText(violation_img, line, (20, y_pos),
                    cv2.FONT_HERSHEY_SIMPLEX, font_scale, color, thickness)

    # 生成文件名并保存
    base_name = os.path.splitext(filename)[0]
    timestamp = time.strftime("%Y%m%d_%H%M%S")
    if is_video:
        violation_filename = f"video_violation_{base_name}_frame{frame_num}_{timestamp}.jpg"
    else:
        violation_filename = f"violation_{base_name}_{timestamp}.jpg"

    violation_path = os.path.join(VIOLATION_FOLDER, violation_filename)
    cv2.imwrite(violation_path, violation_img)
    print(f"🚨 压线违规图片已保存: {violation_path}")

    return violation_path


# -----------------------------
# 8. 单帧处理函数（用于图片和视频）
# -----------------------------
def process_single_frame(frame, frame_info=""):
    """
    处理单帧图像，返回处理结果和统计信息
    """
    if frame is None:
        return None, {}

    orig_h, orig_w = frame.shape[:2]
    print(f"处理帧: {frame_info}, 尺寸: {orig_w}x{orig_h}")

    # 车辆检测推理
    resized_det = cv2.resize(frame, (w_det, h_det))
    input_det = np.expand_dims(resized_det.transpose(2, 0, 1), axis=0)
    detections = vehicle_compiled([input_det])[vehicle_output_layer]

    # 增强版道路分割
    seg_class_resized = enhanced_road_segmentation(frame, PREPROCESS_METHOD)

    # 创建道路标记掩码
    lane_marking_mask = (seg_class_resized == 3)

    # 统计信息
    car_count = 0
    violation_count = 0
    has_violation_in_frame = False

    # 处理每个检测到的车辆
    for det_idx, det in enumerate(detections[0, 0]):
        conf = float(det[2])
        if conf > 0.5:
            x_min = int(det[3] * orig_w)
            y_min = int(det[4] * orig_h)
            x_max = int(det[5] * orig_w)
            y_max = int(det[6] * orig_h)

            vehicle_bbox = [x_min, y_min, x_max, y_max]

            # 压线检测
            is_violation, overlap_ratio, overlap_pixels, bottom_bbox, overlap_mask = check_lane_violation_optimized(
                vehicle_bbox, lane_marking_mask, CONFIDENCE_THRESHOLD, SIDE_MARGIN_RATIO
            )

            # 根据检测结果设置标签
            if is_violation:
                violation_count += 1
                has_violation_in_frame = True
                print(f"🚨 车辆 {det_idx + 1} 压线! 重叠比例: {overlap_ratio:.3f}")

                # 保存压线违规图片
                violation_path = save_violation_image(
                    frame, frame_info, vehicle_bbox, bottom_bbox,
                    overlap_mask, violation_count, car_count + 1, overlap_ratio,
                    is_video="Frame" in frame_info, frame_num=det_idx + 1
                )

            car_count += 1

    # 统计信息
    stats = {
        'total_vehicles': car_count,
        'violations': violation_count,
        'has_violation': has_violation_in_frame,
        'violation_ratio': violation_count / max(car_count, 1)
    }

    return frame, stats  # 返回原帧和统计信息


# -----------------------------
# 9. 稳健的视频读取器
# -----------------------------
class RobustVideoReader:
    """
    稳健的视频读取器，处理损坏的视频文件
    """

    def __init__(self, video_path):
        self.video_path = video_path
        self.cap = None
        self.temp_dir = None
        self.repaired_path = None
        self.current_frame = 0
        self.max_retries = 3

    def open(self):
        """打开视频文件"""
        print(f"🎬 尝试打开视频: {self.video_path}")

        # 首先尝试直接打开
        self.cap = cv2.VideoCapture(self.video_path)
        if self.cap.isOpened():
            print("✅ 直接打开视频成功")
            return True

        # 如果直接打开失败，尝试修复
        print("❌ 直接打开失败，尝试修复视频...")
        self.repaired_path, self.temp_dir = self.repair_video()

        if self.repaired_path and os.path.exists(self.repaired_path):
            self.cap = cv2.VideoCapture(self.repaired_path)
            if self.cap.isOpened():
                print("✅ 使用修复后的视频成功")
                return True

        print("❌ 所有打开视频的尝试都失败了")
        return False

    def repair_video(self):
        """使用FFmpeg修复视频文件"""
        print("🔧 尝试修复视频文件...")

        # 创建临时文件
        temp_dir = tempfile.mkdtemp()
        repaired_video_path = os.path.join(temp_dir, "repaired_video.mp4")

        try:
            # 使用FFmpeg修复视频
            ffmpeg_cmd = [
                'ffmpeg',
                '-y',  # 覆盖输出文件
                '-err_detect', 'ignore_err',  # 忽略错误
                '-i', self.video_path,
                '-c', 'copy',  # 流复制，不重新编码
                '-f', 'mp4',
                repaired_video_path
            ]

            # 运行FFmpeg，忽略所有输出
            result = subprocess.run(
                ffmpeg_cmd,
                stdout=subprocess.DEVNULL,
                stderr=subprocess.DEVNULL,
                timeout=30  # 30秒超时
            )

            if result.returncode == 0 and os.path.exists(repaired_video_path):
                print("✅ 视频修复成功")
                return repaired_video_path, temp_dir
            else:
                print("❌ 视频修复失败，尝试重新编码...")
                return self.repair_video_reencode(temp_dir)

        except Exception as e:
            print(f"❌ 视频修复出错: {e}")
            return self.repair_video_reencode(temp_dir)

    def repair_video_reencode(self, temp_dir):
        """重新编码视频文件"""
        try:
            repaired_video_path = os.path.join(temp_dir, "reencoded_video.mp4")

            ffmpeg_cmd = [
                'ffmpeg',
                '-y',
                '-err_detect', 'ignore_err',
                '-i', self.video_path,
                '-c:v', 'libx264',  # 重新编码视频
                '-c:a', 'aac',  # 重新编码音频
                '-crf', '23',  # 中等质量
                '-preset', 'fast',
                repaired_video_path
            ]

            result = subprocess.run(
                ffmpeg_cmd,
                stdout=subprocess.DEVNULL,
                stderr=subprocess.DEVNULL,
                timeout=60  # 60秒超时
            )

            if result.returncode == 0 and os.path.exists(repaired_video_path):
                print("✅ 视频重新编码成功")
                return repaired_video_path, temp_dir
            else:
                print("❌ 视频重新编码失败")
                return None, temp_dir

        except Exception as e:
            print(f"❌ 视频重新编码出错: {e}")
            return None, temp_dir

    def read(self):
        """读取下一帧，带有错误处理"""
        for attempt in range(self.max_retries):
            try:
                if self.cap is None:
                    return False, None

                ret, frame = self.cap.read()

                # 检查帧是否有效
                if ret and frame is not None and frame.size > 0:
                    self.current_frame += 1
                    return True, frame
                else:
                    # 如果读取失败，尝试重新打开视频
                    print(f"⚠️ 读取帧失败 (尝试 {attempt + 1}/{self.max_retries})")
                    if attempt < self.max_retries - 1:
                        time.sleep(0.1)
                        # 重新尝试读取
                        continue
                    else:
                        return False, None

            except Exception as e:
                print(f"⚠️ 读取帧时出错 (尝试 {attempt + 1}/{self.max_retries}): {e}")
                if attempt < self.max_retries - 1:
                    time.sleep(0.1)  # 短暂等待后重试
                    continue
                else:
                    return False, None

        return False, None

    def get_properties(self):
        """获取视频属性"""
        if self.cap is None:
            return 0, 0, 0, 0

        fps = self.cap.get(cv2.CAP_PROP_FPS)
        total_frames = int(self.cap.get(cv2.CAP_PROP_FRAME_COUNT))
        width = int(self.cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        height = int(self.cap.get(cv2.CAP_PROP_FRAME_HEIGHT))

        return fps, total_frames, width, height

    def release(self):
        """释放资源"""
        if self.cap is not None:
            self.cap.release()

        # 清理临时文件
        if self.temp_dir and os.path.exists(self.temp_dir):
            try:
                shutil.rmtree(self.temp_dir)
                print("✅ 临时文件已清理")
            except Exception as e:
                print(f"⚠️ 清理临时文件时出错: {e}")


# -----------------------------
# 10. 图片批量处理函数
# -----------------------------
def process_images():
    """
    批量处理图片
    """
    image_extensions = {'.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.tif'}
    image_files = []
    for root, dirs, files in os.walk(INPUT_FOLDER):
        for file in files:
            if os.path.splitext(file)[1].lower() in image_extensions:
                image_files.append(os.path.join(root, file))

    if not image_files:
        print(f"❌ 没有在 '{INPUT_FOLDER}' 中找到图片！")
        return

    print(f"✅ 发现 {len(image_files)} 张图片，开始批量处理...")

    violation_stats = []

    for i, img_path in enumerate(image_files, 1):
        print(f"\n--- 处理第 {i}/{len(image_files)} 张: {os.path.basename(img_path)} ---")

        # 读取图像
        image = cv2.imread(img_path)
        if image is None:
            print(f"⚠️ 无法加载图像: {img_path}")
            continue

        # 处理单帧
        result, stats = process_single_frame(image, f"Image {i}")

        if result is not None:
            # 保存结果
            base_name = os.path.splitext(os.path.basename(img_path))[0]
            output_filename = f"{base_name}_result.jpg"
            output_path = os.path.join(OUTPUT_FOLDER, output_filename)

            if os.path.exists(output_path):
                timestamp = time.strftime("%Y%m%d_%H%M%S")
                output_filename = f"{base_name}_result_{timestamp}.jpg"
                output_path = os.path.join(OUTPUT_FOLDER, output_filename)

            cv2.imwrite(output_path, result)
            print(f"✅ 结果已保存至: {output_path}")

            # 记录统计信息
            stats['filename'] = os.path.basename(img_path)
            violation_stats.append(stats)

    # 输出统计报告
    print_statistics(violation_stats, "图片")


# -----------------------------
# 11. 增强的视频处理函数（不保存视频）
# -----------------------------
def process_video():
    """
    处理视频文件，每3帧检测一次，增强错误处理，不保存视频
    """
    if not os.path.exists(VIDEO_INPUT):
        print(f"❌ 视频文件不存在: {VIDEO_INPUT}")
        return

    print(f"🎬 开始处理视频: {VIDEO_INPUT}")
    print("⚠️ 使用稳健的视频读取器，跳过损坏的帧...")
    print("⚠️ 注意: 此模式不保存检测视频，仅输出统计结果和违规图片")

    # 创建稳健的视频读取器
    video_reader = RobustVideoReader(VIDEO_INPUT)

    if not video_reader.open():
        print("❌ 无法打开视频文件")
        return

    # 获取视频属性
    fps, total_frames, width, height = video_reader.get_properties()
    print(f"视频信息: {width}x{height}, FPS: {fps:.2f}, 总帧数: {total_frames}")

    frame_count = 0
    processed_frame_count = 0
    violation_stats = []
    start_time = time.time()
    error_count = 0
    max_errors = 100  # 允许更多错误

    print("开始视频处理（每3帧检测一次）...")

    # 主处理循环
    while True:
        try:
            # 读取下一帧
            ret, frame = video_reader.read()

            if not ret:
                print("✅ 视频读取完成")
                break

            frame_count += 1

            # 检查帧是否有效
            if frame is None or frame.size == 0:
                print(f"⚠️ 第 {frame_count} 帧无效，跳过")
                error_count += 1
                continue

            # 每3帧处理一次
            if frame_count % DETECTION_INTERVAL == 0:
                processed_frame_count += 1
                print(f"--- 处理第 {frame_count}/{total_frames} 帧 ({processed_frame_count}次检测) ---")

                # 处理当前帧
                result, stats = process_single_frame(frame, f"Frame {frame_count}")

                if result is not None:
                    stats['frame_num'] = frame_count
                    violation_stats.append(stats)

            # 显示进度
            if frame_count % 30 == 0:
                elapsed_time = time.time() - start_time
                progress = (frame_count / total_frames) * 100 if total_frames > 0 else 0
                print(f"进度: {progress:.1f}% | 已处理: {elapsed_time:.1f}s")

            # 检查错误限制
            if error_count > max_errors:
                print(f"❌ 错误次数超过限制 ({max_errors})，停止处理")
                break

        except KeyboardInterrupt:
            print("⏹️ 用户中断处理")
            break
        except Exception as e:
            error_count += 1
            print(f"⚠️ 处理第 {frame_count} 帧时发生错误: {e}")
            print(f"    错误类型: {type(e).__name__}")

            if error_count > max_errors:
                print(f"❌ 错误次数超过限制，停止处理")
                break

            continue

    # 释放资源
    video_reader.release()
    cv2.destroyAllWindows()

    # 计算处理时间
    total_time = time.time() - start_time
    processing_fps = frame_count / total_time if total_time > 0 else 0

    print(f"\n✅ 视频处理完成!")
    print(f"总帧数: {frame_count}, 处理帧数: {processed_frame_count}")
    print(f"总耗时: {total_time:.2f}秒, 处理速度: {processing_fps:.2f} FPS")
    print(f"错误帧数: {error_count}")

    # 输出统计报告
    print_statistics(violation_stats, "视频")


# -----------------------------
# 12. 统计报告函数
# -----------------------------
def print_statistics(stats, data_type):
    """
    输出统计报告
    """
    if not stats:
        print(f"❌ 没有{data_type}处理数据！")
        return

    print(f"\n{'=' * 50}")
    print(f"🚗 {data_type}压线检测统计报告")
    print(f"{'=' * 50}")

    total_vehicles = sum(stat['total_vehicles'] for stat in stats)
    total_violations = sum(stat['violations'] for stat in stats)
    total_violation_frames = sum(1 for stat in stats if stat['has_violation'])

    print(f"总处理帧数: {len(stats)}")
    print(f"总检测车辆数: {total_vehicles}")
    print(f"总压线违规数: {total_violations}")
    print(f"总体违规比例: {total_violations / max(total_vehicles, 1) * 100:.1f}%")
    print(f"检测到违规的帧数: {total_violation_frames}")

    if data_type == "视频":
        print(f"检测间隔: 每{DETECTION_INTERVAL}帧检测一次")
        print(f"预处理方法: {PREPROCESS_METHOD}")


# -----------------------------
# 13. 主函数
# -----------------------------
def main():
    """
    主函数 - 选择处理模式
    """
    print("=" * 60)
    print("🚗 车辆压线检测系统（图片+视频）")
    print("=" * 60)
    print("请选择处理模式:")
    print("1. 处理图片文件夹")
    print("2. 处理视频文件（仅检测，不保存视频）")
    print("3. 同时处理图片和视频")

    choice = input("请输入选择 (1/2/3): ").strip()

    if choice == "1":
        print("\n🎯 选择模式: 处理图片文件夹")
        process_images()
    elif choice == "2":
        print("\n🎯 选择模式: 处理视频文件（仅检测，不保存视频）")
        process_video()
    elif choice == "3":
        print("\n🎯 选择模式: 同时处理图片和视频")
        process_images()
        process_video()
    else:
        print("❌ 无效选择，请重新运行程序并输入1、2或3")
        return

    print(f"\n🎉 所有处理完成！")
    print(f"📁 结果目录:")
    print(f"   - 图片结果: {OUTPUT_FOLDER}")
    print(f"   - 违规图片: {VIOLATION_FOLDER}")


# -----------------------------
# 14. 运行主函数
# -----------------------------
if __name__ == "__main__":
    main()